import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from models.TransMIL.TransMIL import TransLayer
from torchinfo import summary

class PPEG_v2(nn.Module):
    def __init__(self, dim=512):
        super().__init__()
        in_dim = dim // 4
        self.coord_fc = nn.Sequential(
            nn.Linear(2, in_dim),
            nn.ReLU(),
            nn.Linear(in_dim, dim),
        )
        
    def _coord_normalize(self, c):
        c = c.float()

        min_c, _ = torch.min(c, dim=0)
        max_c, _ = torch.max(c, dim=0)
        norm_c = (c - min_c) / (max_c - min_c)

        return norm_c
    
    def forward(self, x, coords):
        B, _, C = x.shape
        cls_token, feat_token = x[:, 0], x[:, 1:]

        coords = self._coord_normalize(coords)
        
        pe = self.coord_fc(coords)
        feat_token = feat_token + pe.unsqueeze(0)
        
        x = torch.cat((cls_token.unsqueeze(1), feat_token), dim=1)
        return x


class TransMIL_v2(nn.Module):
    def __init__(self, n_classes):
        super(TransMIL_v2, self).__init__()
        self.pos_layer = PPEG_v2(dim=512)
        self._fc1 = nn.Sequential(nn.Linear(1024, 512), nn.ReLU())
        self.cls_token = nn.Parameter(torch.randn(1, 1, 512))
        self.n_classes = n_classes
        self.layer1 = TransLayer(dim=512)
        self.layer2 = TransLayer(dim=512)
        self.norm = nn.LayerNorm(512)
        self._fc2 = nn.Linear(512, self.n_classes)

    def forward(self, h, coords=None, **kwargs):
        h = h.unsqueeze(0).float() #[1, n, 1024]
        
        h = self._fc1(h) #[B, n, 512]
        
        #---->pad
        # H = h.shape[1]
        # _H, _W = int(np.ceil(np.sqrt(H))), int(np.ceil(np.sqrt(H)))
        # add_length = _H * _W - H
        # h = torch.cat([h, h[:,:add_length,:]],dim = 1) #[B, N, 512]

        #---->cls_token
        B = h.shape[0]
        cls_tokens = self.cls_token.expand(B, -1, -1).cuda()
        h = torch.cat((cls_tokens, h), dim=1)

        #---->Translayer x1
        h = self.layer1(h) #[B, N, 512]

        #---->PPEG
        h = self.pos_layer(h, coords) #[B, N, 512]
        
        #---->Translayer x2
        h = self.layer2(h) #[B, N, 512]

        #---->cls_token
        h = self.norm(h)[:,0]

        #---->predict
        logits = self._fc2(h) #[B, n_classes]
        Y_hat = torch.argmax(logits, dim=1)
        Y_prob = F.softmax(logits, dim = 1)
        
        attn = []
        results_dict = {}
        return logits, Y_prob, Y_hat, attn, results_dict
    
    def relocate(self):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        self.pos_layer = self.pos_layer.to(device)
        self._fc1 = self._fc1.to(device)

        self.layer1 = self.layer1.to(device)
        self.layer2 = self.layer2.to(device)
        self.norm = self.norm.to(device)
        self._fc2 = self._fc2.to(device)



if __name__ == "__main__":

    model = TransMIL_v2(n_classes=2)
    summary(model, ((100, 1024),(100, 2)),depth=3)